# import numpy as np
# import netCDF4 as nc
# import matplotlib.pyplot as plt
# import cartopy.crs as ccrs
# import cartopy.feature as cfeature
# import os
# import glob
# import gc
#
# # 设置路径
# data_dir = "/mnt/datastore/liudddata/result/20200104new"
# output_dir = "./results"
# os.makedirs(output_dir, exist_ok=True)
#
# # 获取文件列表
# nc_files = sorted(glob.glob(os.path.join(data_dir, "*_predicted_2d.nc")))
# print(f"找到 {len(nc_files)} 个数据文件")
#
# # 初始化累加器
# total_cbh = None
# count = None
#
# for i, file in enumerate(nc_files):
#     try:
#         with nc.Dataset(file, 'r') as nf:
#             cbh = nf['predicted'][:].astype(np.float32)
#             cbh = np.ma.filled(cbh, np.nan)
#
#             if total_cbh is None:
#                 total_cbh = np.zeros_like(cbh, dtype=np.float32)
#                 count = np.zeros_like(cbh, dtype=np.int32)
#
#             valid_mask = ~np.isnan(cbh)
#             total_cbh += np.where(valid_mask, cbh, 0)
#             count += valid_mask.astype(np.int32)
#
#         print(f"已处理 {i + 1}/{len(nc_files)} 文件: {os.path.basename(file)}")
#         del cbh
#         gc.collect()
#     except Exception as e:
#         print(f"错误: {file} - {str(e)}")
#         continue
#
# # 计算均值
# mean_cbh = total_cbh / np.where(count > 0, count, np.nan)
#
# # 降采样
# ds_factor = 4
# mean_cbh_lowres = mean_cbh[::ds_factor, ::ds_factor]
# lon = nc.Dataset(nc_files[0])['lon'][:][::ds_factor, ::ds_factor]
# lat = nc.Dataset(nc_files[0])['lat'][:][::ds_factor, ::ds_factor]
#
# # 创建地图画布，使用正交投影
# fig = plt.figure(figsize=(12, 8))
# ax = fig.add_subplot(1, 1, 1, projection=ccrs.Orthographic(central_latitude=0, central_longitude=104.7))
#
# # 添加地理特征
# ax.add_feature(cfeature.COASTLINE)
# ax.add_feature(cfeature.BORDERS, linestyle=':')
# # 将陆地颜色设为白色
# ax.add_feature(cfeature.LAND, facecolor='white')
# ax.add_feature(cfeature.OCEAN, facecolor='white')
#
# # 添加经纬度线
# parallels = range(-90, 91, 30)  # 每30度一个纬线
# meridians = range(-180, 181, 30)  # 每30度一个经线
# gl = ax.gridlines(draw_labels=True, color='gray', linestyle='--', xlocs=meridians, ylocs=parallels)
# gl.top_labels = False
# gl.right_labels = False
# gl.xlabel_style = {'size': 10}
# gl.ylabel_style = {'size': 10}
#
# # 绘图
# cf = ax.contourf(lon, lat, mean_cbh_lowres,
#                  levels=15, cmap='Blues',
#                  transform=ccrs.PlateCarree())
#
# plt.colorbar(cf, label='Cloud Base Height (m)')
# plt.title('2020-1-4 months averaged cbh')
#
# plt.savefig(os.path.join(output_dir, 'average_cbh_optimized.png'), dpi=500, bbox_inches='tight')
# plt.show()

# import numpy as np
# import netCDF4 as nc
# import os
# import glob
# from tqdm import tqdm
#
#
# # 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/season/3_5"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_3_5.nc"
# VARIABLE_NAME = "predicted"
# FILL_VALUE = -999.9
#
#
# def process_2d_grid():
#     # 获取文件列表
#     file_list = sorted(glob.glob(os.path.join(INPUT_DIR, "*_predicted_2d.nc")))
#     if not file_list:
#         raise ValueError("未找到输入文件")
#
#     # 初始化网格存储
#     grid_shape = None
#     lat_2d = None
#     lon_2d = None
#
#     # 读取第一个文件获取网格信息
#     with nc.Dataset(file_list[0], 'r') as template_ds:
#         # 确认原始网格维度（假设维度名为y,x）
#         # if 'y' in template_ds.dimensions and 'x' in template_dimensions:
#         #     y_dim = len(template_ds.dimensions['y'])
#         #     x_dim = len(template_ds.dimensions['x'])
#         #     grid_shape = (y_dim, x_dim)
#         # else:
#         #     raise RuntimeError("原始文件维度不符合预期结构")
#         if 'y' in template_ds.dimensions and 'x' in template_ds.dimensions:
#             y_dim = len(template_ds.dimensions['y'])
#             x_dim = len(template_ds.dimensions['x'])
#             grid_shape = (y_dim, x_dim)
#         else:
#             raise RuntimeError("原始文件维度不符合预期结构")
#
#         # 读取二维坐标变量
#         lat_2d = template_ds['lat'][:].data  # 形状应为(y_dim, x_dim)
#         lon_2d = template_ds['lon'][:].data
#
#     # 初始化数据立方体（时间, y, x）
#     data_cube = np.full((len(file_list), *grid_shape), np.nan, dtype=np.float32)
#     # 使用 np.memmap 创建数据立方体（时间, y, x）
#     # data_cube = np.memmap('temp.dat', dtype=np.float32,
#     #                       mode='w+', shape=(len(file_list), y_dim, x_dim))
#
#     # 读取所有数据
#     for i, file_path in enumerate(tqdm(file_list, desc="处理文件")):
#         with nc.Dataset(file_path, 'r') as ds:
#             # 验证网格一致性
#             if (ds['lat'][:].shape != grid_shape) or (ds['lon'][:].shape != grid_shape):
#                 raise ValueError(f"文件 {file_path} 的网格维度不匹配")
#
#             # 读取并处理数据
#             raw_data = ds[VARIABLE_NAME][:]
#             valid_data = np.ma.filled(raw_data.astype(np.float32), FILL_VALUE)
#             data_cube[i] = np.where(valid_data == FILL_VALUE, np.nan, valid_data)
#
#     # 计算时间平均（忽略nan）
#     mean_cbh = np.nanmean(data_cube, axis=0)
#
#     # 创建输出文件
#     with nc.Dataset(OUTPUT_NC, 'w', format='NETCDF4') as out_ds:
#         # 创建与原始文件相同的维度
#         y_dim = out_ds.createDimension('y', grid_shape[0])
#         x_dim = out_ds.createDimension('x', grid_shape[1])
#
#         # 创建二维坐标变量
#         lat_var = out_ds.createVariable('lat', 'f4', ('y', 'x'))
#         lon_var = out_ds.createVariable('lon', 'f4', ('y', 'x'))
#         lat_var[:] = lat_2d
#         lon_var[:] = lon_2d
#
#         # 添加坐标属性（保持与原始文件一致）
#         lat_var.units = "degrees_north"
#         lat_var.long_name = "latitude"
#         lon_var.units = "degrees_east"
#         lon_var.long_name = "longitude"
#
#         # 创建数据变量（保持原始存储顺序）
#         cbh_var = out_ds.createVariable(
#             'mean_cbh', 'f4', ('y', 'x'),
#             fill_value=FILL_VALUE,
#             zlib=True,
#             complevel=4  # 压缩等级
#         )
#
#         # 处理缺失值并写入
#         output_data = np.where(np.isnan(mean_cbh), FILL_VALUE, mean_cbh)
#         cbh_var[:] = output_data.astype(np.float32)
#
#         # 变量属性
#         cbh_var.units = "m"
#         cbh_var.long_name = "Mean Cloud Base Height"
#         cbh_var.grid_mapping = "fy网格坐标系"  # 根据实际情况填写
#
#         # 全局属性
#         out_ds.title = "二维网格平均云底高度"
#         out_ds.original_grid = f"{grid_shape[0]}x{grid_shape[1]}"
#         out_ds.source_files = "\n".join([os.path.basename(f) for f in file_list])
#
#
# if __name__ == "__main__":
#     try:
#         process_2d_grid()
#         print(f"成功生成二维网格平均数据文件：{OUTPUT_NC}")
#         print(f"输出网格维度验证：{np.array(nc.Dataset(OUTPUT_NC)['mean_cbh'][:]).shape}")
#     except Exception as e:
#         print(f"处理失败：{str(e)}")

import numpy as np
import netCDF4 as nc
import os
import glob
from tqdm import tqdm

# 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/season/3_5"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_3_5.nc"
# 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/season/6_8"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_6_8.nc"
# # 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/season/9_11"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_9_11.nc"
# 配置参数
# INPUT_DIR = "/mnt/datastore/liudddata/season/12_2"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_12_2.nc"
INPUT_DIR = "/mnt/datastore/liudddata/result/20190212_01new"
OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_annual.nc"
VARIABLE_NAME = "predicted"
FILL_VALUE = -999.9


def process_2d_grid():
    # 获取文件列表
    file_list = sorted(glob.glob(os.path.join(INPUT_DIR, "*_predicted_2d.nc")))
    if not file_list:
        raise ValueError("未找到输入文件")

    # 读取第一个文件获取网格信息
    with nc.Dataset(file_list[0], 'r') as template_ds:
        # 验证维度结构
        if 'y' not in template_ds.dimensions or 'x' not in template_ds.dimensions:
            raise RuntimeError("原始文件维度不符合预期结构")
        y_dim = len(template_ds.dimensions['y'])
        x_dim = len(template_ds.dimensions['x'])

        # 获取坐标数据
        lat_2d = template_ds['lat'][:].data
        lon_2d = template_ds['lon'][:].data

    # 初始化累加数组
    sum_data = np.zeros((y_dim, x_dim), dtype=np.float64)  # 使用双精度减少累积误差
    count_data = np.zeros((y_dim, x_dim), dtype=np.uint32)  # 足够存储2173次计数

    # 逐文件处理
    for file_path in tqdm(file_list, desc="处理进度"):
        with nc.Dataset(file_path, 'r') as ds:
            # 快速维度验证
            if ds[VARIABLE_NAME].shape != (y_dim, x_dim):
                continue  # 跳过无效文件

            # 读取并转换数据
            data = ds[VARIABLE_NAME][:].astype(np.float32)

            # 处理填充值
            valid_mask = (data != FILL_VALUE) & (~np.isnan(data))
            valid_data = np.where(valid_mask, data, 0)

            # 更新累加器
            np.add.at(sum_data, (slice(None), slice(None)), valid_data)
            count_data += valid_mask.astype(np.uint32)

    # 计算平均值
    with np.errstate(divide='ignore', invalid='ignore'):
        mean_data = sum_data / count_data

    # 处理无数据区域
    mean_data = np.where(count_data > 0, mean_data, FILL_VALUE).astype(np.float32)

    # 创建输出文件
    with nc.Dataset(OUTPUT_NC, 'w', format='NETCDF4') as out_ds:
        # 定义维度
        out_ds.createDimension('y', y_dim)
        out_ds.createDimension('x', x_dim)

        # 坐标变量
        lat_var = out_ds.createVariable('lat', 'f4', ('y', 'x'), zlib=True)
        lon_var = out_ds.createVariable('lon', 'f4', ('y', 'x'), zlib=True)
        lat_var[:] = lat_2d
        lon_var[:] = lon_2d

        # 数据变量
        cbh_var = out_ds.createVariable(
            'mean_cbh', 'f4', ('y', 'x'),
            fill_value=FILL_VALUE,
            zlib=True,
            complevel=2
        )
        cbh_var[:] = mean_data

        # 元数据属性
        lat_var.units = "degrees_north"
        lon_var.units = "degrees_east"
        cbh_var.long_name = "Average Cloud Base Height"
        cbh_var.units = "m"
        out_ds.title = f"{len(file_list)}文件平均云底高度"


if __name__ == "__main__":
    try:
        process_2d_grid()
        print(f"处理完成，结果保存至：{OUTPUT_NC}")
    except Exception as e:
        print(f"处理失败：{str(e)}")



